# -*- coding: utf-8 -*-

import pandas as pd
from sklearn.model_selection import StratifiedKFold
from otherpkg.utils import timeit


@timeit
def main():
    base_dir = './input'
    train_dir = base_dir + '/train_preliminary'
    train_click_log = pd.read_csv(train_dir + '/click_log.csv')
    train_user = pd.read_csv(train_dir + '/user.csv')
    train_user['label'] = train_user[['age', 'gender']].apply(lambda x: "_".join(map(str, x)), axis=1)
    # 剔除异常用户user_id=839368
    train_click_log.drop(train_click_log[train_click_log.user_id==839368].index, inplace=True)
    for feat in train_click_log.columns:
        train_click_log[feat] = train_click_log[feat].apply(lambda x: None if x == '\\N' else x)

    stratified_label = pd.merge(train_click_log['user_id'], train_user, 'left')
    splited_data_list = []
    n_splits = 3
    skf = StratifiedKFold(n_splits=n_splits)
    for i, (train_idx,valid_idx) in enumerate(skf.split(train_click_log, stratified_label['label'])):
        stratified_train_data, stratified_valid_data = train_click_log.iloc[train_idx], train_click_log.iloc[valid_idx]
        splited_data_list.append((stratified_train_data, stratified_valid_data))

        # save data
        stratified_train_data.to_csv(base_dir + '/train_{}.csv'.format(str(i)), index=False)
        stratified_valid_data.to_csv(base_dir + '/valid_{}.csv'.format(str(i)), index=False)


if __name__ == "__main__":
    main()
